package cas.ihep.alluxio.impl;

import alluxio.Configuration;
import alluxio.PropertyKey;
import alluxio.client.file.options.OutStreamOptions;
import alluxio.client.netty.NettyRPC;
import alluxio.client.netty.NettyRPCContext;
import alluxio.proto.dataserver.Protocol;
import alluxio.util.proto.ProtoMessage;
import alluxio.wire.WorkerNetAddress;
import cas.ihep.alluxio.AlluxioFile;
import com.google.common.base.Preconditions;

import java.io.File;
import java.io.IOException;

public class AlluxioWritableBlock extends AlluxioBaseBlock {

    private static final long FILE_BUFFER_BYTES =
            Configuration.getBytes(PropertyKey.USER_FILE_BUFFER_BYTES);
    private static final long WRITE_TIMEOUT_MS =
            Configuration.getMs(PropertyKey.USER_NETWORK_NETTY_TIMEOUT_MS);

    private NettyRPCContext mNettyRPCContext;
    private ProtoMessage mCreateRequest;
    private long mPosReserved;

    AlluxioWritableBlock(AlluxioFile file, WorkerNetAddress address,
                         long id_, long sz, OutStreamOptions options)throws IOException {
        super(id_, sz);
        AlluxioBaseSystem baseSystem=(AlluxioBaseSystem) file.getSystem();
        AlluxioReadableBlock.ClosableResource resource=mCloser.register(new AlluxioReadableBlock.ClosableResource(
                baseSystem.getContext(),address
        ));
        mCreateRequest = new ProtoMessage(
                Protocol.LocalBlockCreateRequest.newBuilder().setBlockId(id_)
                        .setTier(options.getWriteTier()).setSpaceToReserve(FILE_BUFFER_BYTES).build());
        mNettyRPCContext =
                NettyRPCContext.defaults().setChannel(resource.channel).setTimeout(WRITE_TIMEOUT_MS);
        ProtoMessage message = NettyRPC.call(mNettyRPCContext, mCreateRequest);
        Preconditions.checkState(message.isLocalBlockCreateResponse());
        mPath = message.asLocalBlockCreateResponse().getPath();
    }

    @Override
    public final boolean canRead() {
        return false;
    }

    @Override
    public final boolean canWrite() {
        return true;
    }

    public void cancel() throws IOException {
        if (mClosed) {
            return;
        }
        mClosed = true;
        try {
            NettyRPC.call(mNettyRPCContext, new ProtoMessage(
                    Protocol.LocalBlockCompleteRequest.newBuilder().setBlockId(mId).setCancel(true)
                            .build()));
        } catch (Exception e) {
            throw mCloser.rethrow(e);
        } finally {
            mCloser.close();
        }
    }

    public void close() throws IOException {
        if (mClosed) {
            return;
        }
        File file=new File(mPath);
        mSize=file.length();
        mClosed = true;
        try {
            Protocol.LocalBlockCompleteRequest request =
                    Protocol.LocalBlockCompleteRequest.newBuilder().setBlockId(mId).build();
            NettyRPC.call(mNettyRPCContext, new ProtoMessage(request));
        }finally {
            mCloser.close();
        }
    }

    public void require(long pos) throws IOException {
        if (pos <= mPosReserved) {
            return;
        }
        long toReserve = Math.max(pos - mPosReserved, FILE_BUFFER_BYTES);
        NettyRPC.call(mNettyRPCContext, new ProtoMessage(
                mCreateRequest.asLocalBlockCreateRequest().toBuilder().setSpaceToReserve(toReserve)
                        .setOnlyReserveSpace(true).build()));
        mPosReserved += toReserve;
    }
}
